#!/usr/bin/env bash
# 每天都开心

# 定义颜色
Green="\033[32m"
Font="\033[0m"
Red="\033[31m"

# 检查是否以 root 用户运行
root_need() {
    if [[ $EUID -ne 0 ]]; then
        echo -e "${Red}Error: This script must be run as root!${Font}"
        exit 1
    fi
}

# 检测是否为 OpenVZ 环境
ovz_no() {
    if [[ -d "/proc/vz" ]]; then
        echo -e "${Red}Your VPS is based on OpenVZ, not supported!${Font}"
        exit 1
    fi
}

# 添加 swap
add_swap() {
    echo -e "${Green}请输入需要添加的 swap，建议为内存的 2 倍！${Font}"
    while true; do
        read -p "请输入 swap 数值（单位 MB）: " swapsize
        # 验证输入为数字
        if [[ $swapsize =~ ^[0-9]+$ ]] && [[ $swapsize -gt 0 ]]; then
            break
        else
            echo -e "${Red}请输入有效的数字！${Font}"
        fi
    done

    # 检查是否存在 swapfile
    grep -q "swapfile" /etc/fstab

    # 如果不存在则创建 swap
    if [[ $? -ne 0 ]]; then
        echo -e "${Green}swapfile 未发现，正在为其创建 swapfile...${Font}"
        fallocate -l "${swapsize}M" /swapfile
        chmod 600 /swapfile
        mkswap /swapfile
        swapon /swapfile
        echo '/swapfile none swap defaults 0 0' >> /etc/fstab
        echo -e "${Green}swap 创建成功！以下是信息：${Font}"
        cat /proc/swaps
        cat /proc/meminfo | grep Swap
    else
        echo -e "${Red}swapfile 已存在，swap 设置失败！请先运行脚本删除 swap 后重新设置！${Font}"
    fi
}

# 删除 swap
del_swap() {
    # 检查是否存在 swapfile
    grep -q "swapfile" /etc/fstab

    # 如果存在则移除
    if [[ $? -eq 0 ]]; then
        echo -e "${Green}swapfile 已发现，正在将其移除...${Font}"
        sed -i '/swapfile/d' /etc/fstab
        echo "3" > /proc/sys/vm/drop_caches
        swapoff -a
        rm -f /swapfile
        echo -e "${Green}swap 已删除！${Font}"
    else
        echo -e "${Red}swapfile 未发现，swap 删除失败！${Font}"
    fi
}

# 主菜单
main_menu() {
    echo -e "———————————————————————————————————————"
    echo -e "${Green}Linux VPS 一键添加/删除 swap 脚本${Font}"
    echo -e "${Green}1. 添加 swap${Font}"
    echo -e "${Green}2. 删除 swap${Font}"
    echo -e "${Green}3. 退出脚本${Font}"
    echo -e "———————————————————————————————————————"

    while true; do
        read -p "请输入数字 [1-3]: " num
        case "$num" in
            1)
                add_swap
                break
                ;;
            2)
                del_swap
                break
                ;;
            3)
                echo -e "${Green}退出脚本。${Font}"
                exit 0
                ;;
            *)
                echo -e "${Red}请输入有效的选项 [1-3]！${Font}"
                ;;
        esac
    done
}

# 脚本入口
main() {
    root_need
    ovz_no
    clear
    main_menu
}
main
